package org.zhq.impl;

import org.zhq.core.IoProvider;
import org.zhq.thread.NamedThreadFactory;
import org.zhq.util.CloseUtils;

import java.io.IOException;
import java.nio.channels.*;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * ioProvider的实现
 *
 * @author zhengquan
 * @date 2019/1/24
 */
public class IoSelectorProvider implements IoProvider {
    private final AtomicBoolean closed = new AtomicBoolean(false);

    private final AtomicBoolean inRegInput = new AtomicBoolean(false);
    private final AtomicBoolean inRegOutput = new AtomicBoolean(false);

    private final Selector readSelector;
    private final Selector writeSelector;

    private final HashMap<SelectionKey, Runnable> inputCallbackMap = new HashMap<>();
    private final HashMap<SelectionKey, Runnable> outputCallbackMap = new HashMap<>();

    private final ExecutorService inputHandlePool;
    private final ExecutorService outputHandlePool;

    public IoSelectorProvider() throws IOException {
        //开启selector
        readSelector = Selector.open();
        writeSelector = Selector.open();

        //新建一个4个固定线程的读线程池和一个4个固定线程的写线程池
        inputHandlePool = Executors.newFixedThreadPool(4, new NamedThreadFactory("Io-Input-Thread-"));
        outputHandlePool = Executors.newFixedThreadPool(4, new NamedThreadFactory("Io-Output-Thread-"));

        //开始输入输出
        startRead();
        startWrite();
    }

    private void startRead() {
        Thread thread = new SelectThread("Clink-IoSelectorProvider-Read-Thread", closed, inRegInput,
                readSelector, inputCallbackMap, inputHandlePool, SelectionKey.OP_READ);
        thread.start();
    }


    private void startWrite() {
        Thread thread = new SelectThread("Clink-IoSelectorProvider-Write-Thread", closed, inRegOutput,
                writeSelector, outputCallbackMap, outputHandlePool, SelectionKey.OP_WRITE);
        thread.start();
    }

    @Override
    public void close() throws IOException {
        if (closed.compareAndSet(false, true)) {
            inputHandlePool.shutdown();
            outputHandlePool.shutdown();
            inputCallbackMap.clear();
            outputCallbackMap.clear();
            CloseUtils.close(readSelector, writeSelector);
        }
    }

    private void unRegisterSelection(SocketChannel channel, Selector selector, Map<SelectionKey, Runnable> map, AtomicBoolean locker) {
        synchronized (locker) {
            locker.set(true);
            selector.wakeup();
            try {
                if (channel.isRegistered()) {
                    SelectionKey key = channel.keyFor(selector);
                    if (key != null) {
                        key.cancel();
                        map.remove(key);
                    }
                }
            } finally {
                locker.set(false);
                try {
                    locker.notifyAll();
                } catch (Exception ignored) {
                }
            }
        }
    }

    private SelectionKey registerSelection(SocketChannel channel, Selector selector,
                                           int registerOps, AtomicBoolean locker,
                                           Map<SelectionKey, Runnable> map,
                                           Runnable runnable) {
        synchronized (locker) {
            locker.set(true);
            try {
                selector.wakeup();
                SelectionKey key = null;
                if (channel.isRegistered()) {
                    //查询是否注册过
                    key = channel.keyFor(selector);
                    if (key != null) {
                        key.interestOps(key.readyOps() | registerOps);
                    }
                }

                if (key == null) {
                    //注册selector得到key
                    key = channel.register(selector, registerOps);
                    map.put(key, runnable);
                }
                return key;
            } catch (ClosedChannelException
                    | CancelledKeyException
                    | ClosedSelectorException e) {
                return null;
            } finally {
                locker.set(false);
                try {
                    // 通知
                    locker.notify();
                } catch (Exception ignore) {
                }
            }
        }

    }

    /**
     * @param locker
     */
    private static void waitSelection(final AtomicBoolean locker) {
        synchronized (locker) {
            if (locker.get()) {
                try {
                    locker.wait();
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    /**
     * 处理Selector
     * <p>
     * ps:
     * | 就是  update 一个事件
     * &~  就是  取消 一个事件
     *
     * @param key
     * @param keyOps
     * @param map
     * @param pool
     */
    private static void handleSelection(SelectionKey key, int keyOps, HashMap<SelectionKey, Runnable> map, ExecutorService pool, AtomicBoolean locker) {
        synchronized (locker) {
            try {
                //取消待处理事件的监听
                key.interestOps(key.readyOps() & ~keyOps);
            } catch (CancelledKeyException e) {
                return;
            }
        }
        //得到key对应的runnable
        Runnable runnable = null;
        try {
            runnable = map.get(key);
        } catch (Exception ignored) {

        }
        if (runnable != null && !pool.isShutdown()) {
            pool.execute(runnable);
        }
    }

    @Override
    public void register(HandleProviderCallback callback) throws Exception {
        SelectionKey key;
        if (SelectionKey.OP_READ == callback.ops) {
            key = registerSelection(callback.channel, readSelector, SelectionKey.OP_READ, inRegInput, inputCallbackMap, callback);
        } else {
            key = registerSelection(callback.channel, writeSelector, SelectionKey.OP_WRITE, inRegOutput, outputCallbackMap, callback);
        }

        if (key == null) {
            throw new IOException("注册失败。[通道]:" + callback.channel + "[ops]:" + callback.ops);
        }
    }

    @Override
    public void unregister(SocketChannel channel) {
        unRegisterSelection(channel, readSelector, inputCallbackMap, inRegInput);
        unRegisterSelection(channel, writeSelector, outputCallbackMap, inRegOutput);
    }

    static class SelectThread extends Thread {
        private final AtomicBoolean closed;
        private final AtomicBoolean locker;
        private final Selector selector;
        private final HashMap<SelectionKey, Runnable> callMap;
        private final ExecutorService pool;
        private final int keyOps;

        public SelectThread(String name,
                            AtomicBoolean closed,
                            AtomicBoolean locker,
                            Selector selector,
                            HashMap<SelectionKey, Runnable> callMap,
                            ExecutorService pool, int keyOps) {
            super(name);
            this.closed = closed;
            this.locker = locker;
            this.selector = selector;
            this.callMap = callMap;
            this.pool = pool;
            this.keyOps = keyOps;
            this.setPriority(Thread.MAX_PRIORITY);
        }

        @Override
        public void run() {
            super.run();
            AtomicBoolean closed = this.closed;
            AtomicBoolean locker = this.locker;
            Selector selector = this.selector;
            HashMap<SelectionKey, Runnable> callMap = this.callMap;
            ExecutorService pool = this.pool;
            while (!closed.get()) {
                try {
                    if (selector.select() == 0) {
                        waitSelection(locker);
                        continue;
                    } else if (locker.get()) {
                        waitSelection(locker);
                    }
                    Iterator<SelectionKey> iterator = selector.selectedKeys().iterator();
                    selector.selectedKeys().iterator();
                    while (iterator.hasNext()) {
                        SelectionKey key = iterator.next();
                        if (key.isValid()) {
                            handleSelection(key, keyOps, callMap, pool, locker);
                        }
                        iterator.remove();
                    }
                } catch (IOException e) {
                    e.printStackTrace();
                } catch (ClosedSelectorException ignored) {
                    break;
                }
            }
        }
    }
}
